# 切分数据集
# -*- coding: utf-8 -*-
# @Time    : 2023-12-14 16:11
# @Author  : Jiang Liu
import os
import random
from tqdm import tqdm

import shutil


def split_dataset(imgs_dir, txts_dir, train_ratio=0.8, val_ratio=0.1, test_ratio=0.1):
    # 获取所有图片文件
    img_filenames = [f for f in os.listdir(imgs_dir) if f.endswith('.jpg')]

    # 随机打乱图片和标注文件的顺序
    random.seed(2023)
    random.shuffle(img_filenames)

    # 计算训练集、验证集和测试集的数量
    num_imgs = len(img_filenames)
    num_train = int(num_imgs * train_ratio)
    num_val = int(num_imgs * val_ratio)
    num_test = num_imgs - num_train - num_val

    train_img = []
    train_txt = []
    val_img = []
    val_txt = []
    test_img = []
    test_txt = []

    # 将图片和标注文件按照划分结果保存到对应的文件夹中
    for idx, img_filename in tqdm(enumerate(img_filenames)):
        basename = os.path.splitext(img_filename)[0]
        img_filepath = os.path.join(imgs_dir, img_filename)
        txt_filepath = os.path.join(txts_dir, basename + '.txt')
        if idx < num_train:
            train_img.append(img_filepath)
            train_txt.append(txt_filepath)
        elif idx < num_train + num_val:
            val_img.append(img_filepath)
            val_txt.append(txt_filepath)
        else:
            test_img.append(img_filepath)
            test_txt.append(txt_filepath)
    return train_img, train_txt, val_img, val_txt, test_img, test_txt


def copyfiles(filepaths, dst_dir):
    for filepath in tqdm(filepaths):
        shutil.copy(filepath, dst_dir)


def split_result(filelist, subtype):
    prefix = './images/train'
    if subtype == 'val':
        prefix = './images/val'
    elif subtype == 'test':
        prefix = './images/test'
    new_filelist = []
    for filepath in filelist:
        new_filelist.append(prefix + "/" + os.path.basename(filepath))
    return new_filelist


def main():
    imgs_dir = '../processed/images/'
    txts_dir = '../processed/labels/'
    train_img, train_txt, val_img, val_txt, test_img, test_txt = split_dataset(imgs_dir, txts_dir, 0.9, 0.1, 0)
    train_img_dir = '../splited/images/train'
    val_img_dir = '../splited/images/val'
    test_img_dir = '../splited/images/test'
    train_txt_dir = '../splited/labels/train'
    val_txt_dir = '../splited/labels/val'
    test_txt_dir = '../splited/labels/test'
    for dir_path in [train_img_dir, val_img_dir, test_img_dir, train_txt_dir, val_txt_dir, test_txt_dir]:
        os.makedirs(dir_path, exist_ok=True)
    copyfiles(train_img, train_img_dir)
    copyfiles(val_img, val_img_dir)
    copyfiles(test_img, test_img_dir)
    copyfiles(train_txt, train_txt_dir)
    copyfiles(val_txt, val_txt_dir)
    copyfiles(test_txt, test_txt_dir)
    # 将划分结果写入到 train.txt、val.txt 和 test.txt
    with open('../splited/train.txt', 'w') as f:
        f.write('\n'.join(split_result(train_img, 'train')))
    with open('../splited/val.txt', 'w') as f:
        f.write('\n'.join(split_result(val_img, 'val')))
    with open('../splited/test.txt', 'w') as f:
        f.write('\n'.join(split_result(test_img, 'test')))


if __name__ == '__main__':
    main()
